#!/usr/bin/env python
# -*- coding: utf-8 -*-
import rospy, sys
import moveit_commander
import math,random
from std_msgs.msg import Int8,Empty
from xarm_moveit_demo.srv import *
from geometry_msgs.msg import PoseStamped, Pose
from moveit_commander import MoveGroupCommander, PlanningSceneInterface
from moveit_msgs.msg import PlanningScene, ObjectColor
from moveit_msgs.msg import Grasp,PlaceLocation,GripperTranslation, MoveItErrorCodes
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from ar_track_alvar_msgs.msg import AlvarMarkers
from tf.transformations import quaternion_from_euler,euler_from_quaternion
from copy import deepcopy
from xarm_vision.srv import *
from xarm_driver.msg import SingleJointControl

GROUP_NAME_ARM = 'xarm'
GROUP_NAME_GRIPPER = 'gripper'
GRIPPER_FRAME = 'gripper_centor_link'
GRIPPER_OPEN = [0.69,0.69]
GRIPPER_GRASP = [0.2,0.2]
GRIPPER_CLOSED = [0.0,0.0]
GRIPPER_JOINT_NAMES = ['gripper_1_joint','gripper_2_joint']
GRIPPER_EFFORT = [1.0,1.0]
REFERENCE_FRAME = 'base_link'
class XarmCompete:
    def __init__(self):
        # 初始化Python API 依赖的moveit_commanderC++系统
        moveit_commander.roscpp_initialize(sys.argv)
        # 初始化节点
        rospy.init_node('xarm_compete_pickup')
        # 话题ar_pose_marker的订阅端，回调函数为get_tags
        rospy.Subscriber("ar_pose_marker", AlvarMarkers, self.get_tags)
        # 初始化tag_result为AlvarMarkers消息类型
        self.tag_result = AlvarMarkers()

        self.single_joint_control = rospy.Publisher('arm/commands/single_joint_control', SingleJointControl,queue_size=10)

        # 初始化场景对象，用来在规划场景中添加或移除物体
        self.scene = PlanningSceneInterface()
        rospy.sleep(1)

        self.init_flag = 0

        self.table_id = "table"
        # 设置桌子的长宽高 [l, w, h]
        table_size = [1.0, 1.2, 0.01]
        # 设置桌子的位姿
        table_pose = PoseStamped()
        table_pose.header.frame_id = REFERENCE_FRAME
        table_pose.pose.position.x = 0.0
        table_pose.pose.position.y = 0.0
        table_pose.pose.position.z = - table_size[2] / 2.0
        table_pose.pose.orientation.w = 1.0
        # 把桌子添加到规划场景中
        self.scene.add_box(self.table_id, table_pose, table_size)
        rospy.sleep(1)

        self.target_ids = [1,2,3,4]
        self.target_1_pose = PoseStamped()
        #self.target_2_pose = PoseStamped()
        self.target_3_pose = PoseStamped()
        #self.target_4_pose = PoseStamped()
        self.yaw_1_offset = 0


        # Initialize the move group for the xarm
        self.xarm = MoveGroupCommander(GROUP_NAME_ARM)

        # Initialize the move group for the gripper
        self.gripper = MoveGroupCommander(GROUP_NAME_GRIPPER)

        # 机械臂先回到初始位置
        self.xarm.set_named_target('Home')
        self.xarm.go()
        rospy.sleep(1)

        # 定义ROS服务/compete/pick_place的服务端，服务回调函数为call_pick_place()
        self.pick_place_srv = rospy.Service('/compete/pick_place', CallPickup, self.call_pick_place)

        rospy.loginfo(" Pick and Place demo is ready.  ")

        rospy.spin()

        # Shut down MoveIt cleanly
        moveit_commander.roscpp_shutdown()

        # Exit the script
        moveit_commander.os._exit(0)



    # ar_pose_marker话题回调函数
    def get_tags(self, msg):
        self.tag_result = msg
        for tag in self.tag_result.markers:
            tag.pose.header.frame_id = REFERENCE_FRAME
            if tag.id == 1:
                self.target_1_pose = tag.pose
#            elif tag.id == 2:
#                self.target_2_pose = tag.pose
            elif tag.id == 3:
                self.target_3_pose = tag.pose
#            elif tag.id == 4:
#                self.target_4_pose = tag.pose
            else:
                continue


    def init_target_1_scene(self):
        object_1_size = [0.062,0.062,0.07]
        # 获取AR标签的姿态，从四元数转为欧拉角表示
        rpy_euler = euler_from_quaternion([self.target_1_pose.pose.orientation.x, self.target_1_pose.pose.orientation.y,self.target_1_pose.pose.orientation.z,self.target_1_pose.pose.orientation.w ])
        # yaw_offset列表用于保存每个目标方块的偏航角与抓取时的gripper_centor_link的偏航角yaw的差值
        yaw_1 = math.atan(self.target_1_pose.pose.position.y/self.target_1_pose.pose.position.x)
        self.yaw_1_offset = (rpy_euler[2]-yaw_1)
        # 设置方块的姿态与欧拉角的Yaw偏航角一致，Roll和Pitch为零
        q = quaternion_from_euler(0, 0, rpy_euler[2])
        self.target_1_pose.pose.orientation.x = q[0]
        self.target_1_pose.pose.orientation.y = q[1]
        self.target_1_pose.pose.orientation.z = q[2]
        self.target_1_pose.pose.orientation.w = q[3]
        object_1_pose = self.target_1_pose
        object_1_pose.pose.position.z = object_1_size[2]/2
        self.scene.add_box(str(1), object_1_pose, object_1_size)

    def init_target_2_scene(self):
        object_2_size = [0.062,0.062,0.13]
        # 获取AR标签的姿态，从四元数转为欧拉角表示
        rpy_euler = euler_from_quaternion([self.target_2_pose.pose.orientation.x, self.target_2_pose.pose.orientation.y,self.target_2_pose.pose.orientation.z,self.target_2_pose.pose.orientation.w ])
        # yaw_offset列表用于保存每个目标方块的偏航角与抓取时的gripper_centor_link的偏航角yaw的差值
        yaw_2 = math.atan(self.target_2_pose.pose.position.y/self.target_2_pose.pose.position.x)
        self.yaw_2_offset = (rpy_euler[2]-yaw_2)
        # 设置方块的姿态与欧拉角的Yaw偏航角一致，Roll和Pitch为零
        q = quaternion_from_euler(0, 0, rpy_euler[2])
        self.target_2_pose.pose.orientation.x = q[0]
        self.target_2_pose.pose.orientation.y = q[1]
        self.target_2_pose.pose.orientation.z = q[2]
        self.target_2_pose.pose.orientation.w = q[3]
        object_2_pose = self.target_2_pose
        object_2_pose.pose.position.z = object_2_size[2]/2
        self.scene.add_box(str(2), object_2_pose, object_2_size)

        object_3_size = [0.062,0.062]
        object_4_size = [0.062,0.062,0.12]

    #
    def init_target_3_scene(self):
        object_3_size = [0.078,0.078,0.15]
        # 获取AR标签的姿态，从四元数转为欧拉角表示
#        rpy_euler = euler_from_quaternion([self.target_3_pose.pose.orientation.x, self.target_3_pose.pose.orientation.y,self.target_3_pose.pose.orientation.z,self.target_3_pose.pose.orientation.w ])
#        # yaw_offset列表用于保存每个目标方块的偏航角与抓取时的gripper_centor_link的偏航角yaw的差值
        yaw_3 = math.atan(self.target_3_pose.pose.position.y/self.target_3_pose.pose.position.x)
#        self.yaw_3_offset = (rpy_euler[2]-yaw_3)
        self.yaw_3_offset = 0
        # 设置方块的姿态与欧拉角的Yaw偏航角一致，Roll和Pitch为零
        q = quaternion_from_euler(0, 0, yaw_3)
        self.target_3_pose.pose.orientation.x = q[0]
        self.target_3_pose.pose.orientation.y = q[1]
        self.target_3_pose.pose.orientation.z = q[2]
        self.target_3_pose.pose.orientation.w = q[3]
        object_3_pose = self.target_3_pose
        object_3_pose.pose.position.z = object_3_size[2]/2
        self.scene.add_box(str(3), object_3_pose, object_3_size)

    #
    def init_target_4_scene(self):
        object_4_size = [0.05,0.05,0.16]
        # 获取AR标签的姿态，从四元数转为欧拉角表示
        rpy_euler = euler_from_quaternion([self.target_4_pose.pose.orientation.x, self.target_4_pose.pose.orientation.y,self.target_4_pose.pose.orientation.z,self.target_4_pose.pose.orientation.w ])
        # yaw_offset列表用于保存每个目标方块的偏航角与抓取时的gripper_centor_link的偏航角yaw的差值
        yaw_4 = math.atan(self.target_4_pose.pose.position.y/self.target_4_pose.pose.position.x)
        self.yaw_4_offset = (rpy_euler[2]-yaw_4)
        # 设置方块的姿态与欧拉角的Yaw偏航角一致，Roll和Pitch为零
        q = quaternion_from_euler(0, 0, rpy_euler[2])
        self.target_4_pose.pose.orientation.x = q[0]
        self.target_4_pose.pose.orientation.y = q[1]
        self.target_4_pose.pose.orientation.z = q[2]
        self.target_4_pose.pose.orientation.w = q[3]
        self.target_4_pose.pose.position.y -= 0.025

        object_4_pose = self.target_4_pose
        object_4_pose.pose.position.z = object_4_size[2]/2
        self.scene.add_box(str(4), object_4_pose, object_4_size)

    # /xarm_vision_pickup的服务回调函数
    def call_pick_place(self, req):
        rospy.sleep(1)
        self.init_target_1_scene()
        self.init_target_3_scene()
        rospy.sleep(1)
#        if self.init_flag == 0:
#            self.init_target_1_scene()
#            self.init_target_2_scene()
#            self.init_target_3_scene()
#            self.init_target_4_scene()
#            rospy.sleep(1)
#            self.init_flag =1

        # 设置桌子table为抓取和放置操作的支撑面，使MoveIt!忽略物体放到桌子上时产生的碰撞警告
        self.xarm.set_support_surface_name(self.table_id)
        if req.target_id == 0:
            arm_down = [0,-0.4,0,-1.23,0,-0.07]
            self.xarm.set_joint_value_target(arm_down)
            self.xarm.go()
            self.gripper.set_joint_value_target(GRIPPER_OPEN)
            self.gripper.go()
            self.gripper.set_joint_value_target(GRIPPER_CLOSED)
            self.gripper.go()
            # 回到初始位姿
            self.xarm.set_named_target('Home')
            self.xarm.go()

        elif req.target_id == 1:
            self.init_target_1_scene()
            # 设置一个放置目标位姿place_pose
            place_pose = PoseStamped()
            place_pose.header.frame_id = REFERENCE_FRAME
            place_pose.pose.position.x = random.randint(24,30)/100.0
            place_pose.pose.position.y = random.randint(-15,15)/100.0
            place_pose.pose.position.z = 0.035

            # 抓取目标并放置到指定位置
            place_yaw = math.atan(place_pose.pose.position.y/place_pose.pose.position.x)
            q = quaternion_from_euler(0, 0, place_yaw + self.yaw_1_offset)
            place_pose.pose.orientation.x = q[0]
            place_pose.pose.orientation.y = q[1]
            place_pose.pose.orientation.z = q[2]
            place_pose.pose.orientation.w = q[3]
            if self.pick_and_place(req.target_id, self.target_1_pose, place_pose):
                # 回到初始位姿
                self.xarm.set_named_target('Home')
                self.xarm.go()
#                # 闭合手爪
#                self.gripper.set_joint_value_target(GRIPPER_CLOSED)
#                self.gripper.go()
            else:
                return CallPickupResponse(False)

        elif req.target_id == 2:
            self.init_target_2_scene()
            # 设置一个放置目标位姿place_pose
            place_pose = PoseStamped()
            place_pose.header.frame_id = REFERENCE_FRAME
            place_pose.pose.position.x = 0.1
            place_pose.pose.position.y = -0.3
            place_pose.pose.position.z = 0.065

            # 抓取目标并放置到指定位置
            place_yaw = math.atan(place_pose.pose.position.y/place_pose.pose.position.x)
            q = quaternion_from_euler(0, 0, place_yaw + self.yaw_2_offset)
            place_pose.pose.orientation.x = q[0]
            place_pose.pose.orientation.y = q[1]
            place_pose.pose.orientation.z = q[2]
            place_pose.pose.orientation.w = q[3]
            if self.pick_and_place(req.target_id, self.target_2_pose, place_pose):
                # 回到初始位姿
                self.xarm.set_named_target('Home')
                self.xarm.go()
                # 闭合手爪
                self.gripper.set_joint_value_target(GRIPPER_CLOSED)
                self.gripper.go()
            else:
                return CallPickupResponse(False)
        #
        elif req.target_id == 3:
            self.init_target_3_scene()
            # 设置一个放置目标位姿place_pose，区域内某个随机位置
            place_pose = PoseStamped()
            place_pose.header.frame_id = REFERENCE_FRAME
            place_pose.pose.position.x = random.randint(45,46)/100.0
            place_pose.pose.position.y = random.randint(-18,18)/100.0
            place_pose.pose.position.z = 0.075

            # 抓取目标并放置
            place_yaw = math.atan(place_pose.pose.position.y/place_pose.pose.position.x)
            q = quaternion_from_euler(0, 0, place_yaw + self.yaw_3_offset)
            place_pose.pose.orientation.x = q[0]
            place_pose.pose.orientation.y = q[1]
            place_pose.pose.orientation.z = q[2]
            place_pose.pose.orientation.w = q[3]
            if self.pick_and_place(req.target_id, self.target_3_pose, place_pose):
                # 回到初始位姿
                self.xarm.set_named_target('Home')
                self.xarm.go()
#                # 闭合手爪
#                self.gripper.set_joint_value_target(GRIPPER_CLOSED)
                self.gripper.go()
            else:
                return CallPickupResponse(False)
        #
        elif req.target_id == 4:
            self.init_target_4_scene()
            # 设置一个放置目标位姿place_pose
            place_pose = PoseStamped()
            place_pose.header.frame_id = REFERENCE_FRAME
            place_pose.pose.position.x = 0.4
            place_pose.pose.position.y = -0.3
            place_pose.pose.position.z = 0.08

            # 抓取目标并放置到指定位置
            place_yaw = math.atan(place_pose.pose.position.y/place_pose.pose.position.x)
            q = quaternion_from_euler(0, 0, place_yaw + self.yaw_4_offset)
            place_pose.pose.orientation.x = q[0]
            place_pose.pose.orientation.y = q[1]
            place_pose.pose.orientation.z = q[2]
            place_pose.pose.orientation.w = q[3]
            if self.pick_and_place(req.target_id, self.target_4_pose, place_pose):
                # 回到初始位姿
                self.xarm.set_named_target('Home')
                self.xarm.go()
                # 闭合手爪
                self.gripper.set_joint_value_target(GRIPPER_CLOSED)
                self.gripper.go()
            else:
                return CallPickupResponse(False)
        elif req.target_id == 5:
            joint_control = SingleJointControl()
            joint_control.id = 3
            joint_control.rad = 1.57
            self.single_joint_control.publish(joint_control)
            rospy.sleep(2)
            joint_control.id = 4
            joint_control.rad = -0.4
            self.single_joint_control.publish(joint_control)
            rospy.sleep(1)
            joint_control.rad = 0.4
            self.single_joint_control.publish(joint_control)
            rospy.sleep(2)
            joint_control.rad = 0
            self.single_joint_control.publish(joint_control)
            rospy.sleep(1)
            joint_control.id = 3
            joint_control.rad = 0
            self.single_joint_control.publish(joint_control)
            rospy.sleep(1)

        return CallPickupResponse(True)




    def pick_and_place(self,target_id,target_pose,place_pose):
        # 获取 end-effector link的名字
        end_effector_link = self.xarm.get_end_effector_link()

        # 允许重规划
        self.xarm.allow_replanning(True)

        # 设置目标的参考系
        self.xarm.set_pose_reference_frame(REFERENCE_FRAME)

        # 设置每次规划尝试的时间为5s
        self.xarm.set_planning_time(5)

        # 定义抓取操作的最大尝试次数
        max_pick_attempts = 3

        # 定义放置操作的最大尝试次数
        max_place_attempts = 3

        grasp = Grasp()
        # Set the allowed touch objects to the input list
        grasp.allowed_touch_objects = self.table_id

        # Don't restrict contact force
        grasp.max_contact_force = 0

        # 设置 pre_grasp_posture
        grasp.pre_grasp_posture = self.make_gripper_posture(GRIPPER_OPEN)

        # 设置grasp_pose
        grasp.grasp_pose = target_pose
        # 设置抓取的初始姿态为机械臂末端比较容易到达的位姿
        yaw = math.atan(grasp.grasp_pose.pose.position.y/grasp.grasp_pose.pose.position.x)

        if target_id == 1:
            grasp.grasp_posture = self.make_gripper_posture(GRIPPER_GRASP)
            # 设置靠近和撤离的方向以及距离
            grasp.pre_grasp_approach = self.make_gripper_translation(0.07, 0.1, [0,0,-1])
            grasp.post_grasp_retreat = self.make_gripper_translation(0.08, 0.12, [0.0, 0.0, 1.0])
            grasp.grasp_pose.pose.position.z = 0.07
            q = quaternion_from_euler(0, math.pi/2, yaw)
            # Set the grasp pose orientation accordingly
            grasp.grasp_pose.pose.orientation.x = q[0]
            grasp.grasp_pose.pose.orientation.y = q[1]
            grasp.grasp_pose.pose.orientation.z = q[2]
            grasp.grasp_pose.pose.orientation.w = q[3]
            grasp.id = str(1)

        elif target_id == 2:
            grasp.grasp_posture = self.make_gripper_posture(GRIPPER_GRASP)
            # 设置靠近和撤离的方向以及距离
            grasp.pre_grasp_approach = self.make_gripper_translation(0.07, 0.1, [0,0,-1])
            grasp.post_grasp_retreat = self.make_gripper_translation(0.08, 0.12, [0.0, 0.0, 1.0])
            grasp.grasp_pose.pose.position.z = 0.13
            q = quaternion_from_euler(0, math.pi/2, yaw)
            # Set the grasp pose orientation accordingly
            grasp.grasp_pose.pose.orientation.x = q[0]
            grasp.grasp_pose.pose.orientation.y = q[1]
            grasp.grasp_pose.pose.orientation.z = q[2]
            grasp.grasp_pose.pose.orientation.w = q[3]
            grasp.id = str(2)

        elif target_id == 3:
            grasp.grasp_posture = self.make_gripper_posture(GRIPPER_GRASP)
            # 设置靠近和撤离的方向以及距离
            grasp.pre_grasp_approach = self.make_gripper_translation(0.07, 0.2, [grasp.grasp_pose.pose.position.x,grasp.grasp_pose.pose.position.y,0])
            grasp.post_grasp_retreat = self.make_gripper_translation(0.06, 0.15, [0.0, 0.0, 1.0])
            grasp.grasp_pose.pose.position.z = 0.12
            q = quaternion_from_euler(0, 0, yaw)
            # Set the grasp pose orientation accordingly
            grasp.grasp_pose.pose.orientation.x = q[0]
            grasp.grasp_pose.pose.orientation.y = q[1]
            grasp.grasp_pose.pose.orientation.z = q[2]
            grasp.grasp_pose.pose.orientation.w = q[3]
            grasp.id = str(3)
        #
        elif target_id == 4:
            grasp.grasp_posture = self.make_gripper_posture([0.13,0.13])
            # 设置靠近和撤离的方向以及距离
            grasp.pre_grasp_approach = self.make_gripper_translation(0.08, 0.18, [grasp.grasp_pose.pose.position.x,grasp.grasp_pose.pose.position.y,0])
            grasp.post_grasp_retreat = self.make_gripper_translation(0.06, 0.12, [0.0, 0.0, 1.0])
            grasp.grasp_pose.pose.position.z = 0.14
            q = quaternion_from_euler(0, 0, yaw)
            # Set the grasp pose orientation accordingly
            grasp.grasp_pose.pose.orientation.x = q[0]
            grasp.grasp_pose.pose.orientation.y = q[1]
            grasp.grasp_pose.pose.orientation.z = q[2]
            grasp.grasp_pose.pose.orientation.w = q[3]
            grasp.id = str(4)


        # 设置result标记每次抓取尝试的结果
        result = None
        n_attempts = 0
        grasps = []
        grasps.append(deepcopy(grasp))
        # 循环尝试抓取操作直到抓取成功或最大尝试次数用尽
        while result != MoveItErrorCodes.SUCCESS and n_attempts < max_pick_attempts:
            n_attempts += 1
            rospy.loginfo("Pick attempt: " +  str(n_attempts))
            result = self.xarm.pick(str(target_id), grasps)
            rospy.sleep(0.2)

        # 如果抓取成功，尝试物品放置操作。
        if result == MoveItErrorCodes.SUCCESS:
            # 重置标记操作结果的result和尝试次数
            result = None
            n_attempts = 0

            # 生成一系列放置位姿
            places = self.make_places(place_pose)

            # 循环放置直到放置成功或超过最大尝试次数
            while result != MoveItErrorCodes.SUCCESS and n_attempts < max_place_attempts:
                n_attempts += 1
                rospy.loginfo("Place attempt: " +  str(n_attempts))
                for place in places:
                    result = self.xarm.place(str(target_id),place)
                    if result == MoveItErrorCodes.SUCCESS:
                        return True
                rospy.sleep(0.2)

            if result != MoveItErrorCodes.SUCCESS:
                rospy.loginfo("Place operation failed after " + str(n_attempts) + " attempts.")
                return False
        else:
            rospy.loginfo("Pick operation failed after " + str(n_attempts) + " attempts.")
            return False



    # 通过手爪关节的位置生成JointTrajectory消息类型的轨迹返回值
    def make_gripper_posture(self, joint_positions):
        # Initialize the joint trajectory for the gripper joints
        t = JointTrajectory()

        # Set the joint names to the gripper joint names
        t.joint_names = GRIPPER_JOINT_NAMES

        # Initialize a joint trajectory point to represent the goal
        tp = JointTrajectoryPoint()

        # Assign the trajectory joint positions to the input positions
        tp.positions = joint_positions

        # Set the gripper effort
        tp.effort = GRIPPER_EFFORT

        tp.time_from_start = rospy.Duration(5)

        # Append the goal point to the trajectory points
        t.points.append(tp)

        # Return the joint trajectory
        return t

    # 通过传入的向量vector和最小距离、期望距离，生成GripperTranslation消息类型并返回
    def make_gripper_translation(self, min_dist, desired, vector):
        # Initialize the gripper translation object
        g = GripperTranslation()

        # Set the direction vector components to the input
        g.direction.vector.x = vector[0]
        g.direction.vector.y = vector[1]
        g.direction.vector.z = vector[2]

        # The vector is relative to the gripper frame
        g.direction.header.frame_id = REFERENCE_FRAME

        # Assign the min and desired distances from the input
        g.min_distance = min_dist
        g.desired_distance = desired
        return g



    # 生成一系列可能的放置点位姿
    def make_places(self, init_pose):
        # Initialize the place location as a PlaceLocation message
        place = PlaceLocation()
        # Start with the input place pose
        place.place_pose = init_pose
        # 设置靠近放置点的方向、最小移动距离和期望距离，这里设置为沿着Z轴向下移动0.08米
        place.pre_place_approach.direction.header.frame_id = REFERENCE_FRAME;
        place.pre_place_approach.direction.vector.z = -1.0;
        place.pre_place_approach.min_distance = 0.05;
        place.pre_place_approach.desired_distance = 0.08;
        # 设置放置完成后机械臂的撤离方向、移动最小距离和期望距离，这里设置为沿着Z轴向上移动0.08米
        place.post_place_retreat.direction.header.frame_id = REFERENCE_FRAME;
        place.post_place_retreat.direction.vector.z = 1.0;
        place.post_place_retreat.min_distance = 0.08;
        place.post_place_retreat.desired_distance = 0.1;

        # A list of x shifts (meters) to try
        x_vals = [0, 0.005, -0.005]

        # A list of y shifts (meters) to try
        y_vals = [0, 0.005, -0.005]

        # A list to hold the places
        places = []

        # Generate a place pose for each angle and translation
        for y in y_vals:
            for x in x_vals:
                place.place_pose.pose.position.x = init_pose.pose.position.x + x
                place.place_pose.pose.position.y = init_pose.pose.position.y + y
                # Append this place pose to the list
                places.append(deepcopy(place))

        # Return the list
        return places

if __name__ == "__main__":
    XarmCompete()
